import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.linear_model import LinearRegression, Ridge
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import warnings
warnings.filterwarnings('ignore')

class DataPreprocessor:
    """数据预处理类"""
    
    def __init__(self):
        self.scaler = MinMaxScaler()
        self.date_features = None
        
    def create_date_features(self, df, date_col='finish_time'):
        """创建日期特征"""
        df = df.copy()
        df[date_col] = pd.to_datetime(df[date_col])
        
        # 基础时间特征
        df['year'] = df[date_col].dt.year
        df['month'] = df[date_col].dt.month
        df['day'] = df[date_col].dt.day
        df['dayofweek'] = df[date_col].dt.dayofweek
        df['dayofyear'] = df[date_col].dt.dayofyear
        df['weekofyear'] = df[date_col].dt.isocalendar().week
        
        # 周期性特征（使用正弦余弦变换）
        df['month_sin'] = np.sin(2 * np.pi * df['month'] / 12)
        df['month_cos'] = np.cos(2 * np.pi * df['month'] / 12)
        df['day_sin'] = np.sin(2 * np.pi * df['dayofweek'] / 7)
        df['day_cos'] = np.cos(2 * np.pi * df['dayofweek'] / 7)
        
        # 趋势特征
        df['days_from_start'] = (df[date_col] - df[date_col].min()).dt.days
        
        return df
    
    def fill_missing_values(self, df, target_col='value', method='advanced'):
        """填充缺失值"""
        df = df.copy()
        
        if method == 'simple':
            # 简单方法：用0填充
            df[target_col] = df[target_col].fillna(0)
            
        elif method == 'mean':
            # 均值填充
            mean_val = df[target_col].mean()
            df[target_col] = df[target_col].fillna(mean_val if not pd.isna(mean_val) else 0)
            
        elif method == 'advanced':
            # 高级方法：结合时间邻近性和周期性模式
            
            # 首先用0填充所有空值（假设没有销售的日子销售额为0）
            df[target_col] = df[target_col].fillna(0)
            
            # 如果有足够的历史数据，考虑周期性模式
            if len(df) > 7:
                # 计算同星期的平均销售额
                df['dayofweek'] = pd.to_datetime(df['finish_time']).dt.dayofweek
                weekly_avg = df.groupby('dayofweek')[target_col].mean()
                
                # 对于原本为空的值，用同星期的平均值替代（如果有的话）
                mask = df[target_col] == 0  # 假设原本为0的是空值填充的
                for idx in df[mask].index:
                    day_of_week = df.loc[idx, 'dayofweek']
                    if weekly_avg[day_of_week] > 0:
                        df.loc[idx, target_col] = weekly_avg[day_of_week]
        
        return df
    
    def prepare_features(self, df, date_col='finish_time', target_col='value'):
        """完整的特征准备流程"""
        print(f"准备特征数据...")
        
        # 1. 处理日期列
        df_processed = self.create_date_features(df, date_col)
        
        # 2. 处理空值
        if target_col and target_col in df_processed.columns:
            df_processed = self.fill_missing_values(df_processed, target_col)
        
        # 3. 准备特征矩阵 - 确保至少有基本的特征
        feature_cols = ['year', 'month', 'day', 'dayofweek', 'dayofyear', 'weekofyear',
                       'month_sin', 'month_cos', 'day_sin', 'day_cos', 'days_from_start']
        available_features = [col for col in feature_cols if col in df_processed.columns]
        
        if not available_features:
            # 如果没有其他特征，创建基本的时间序列特征
            df_processed['time_index'] = np.arange(len(df_processed))
            available_features = ['time_index']
        
        X = df_processed[available_features]
        
        # 4. 标准化特征
        if len(X) > 1:  # 确保有足够的数据
            X_scaled = self.scaler.fit_transform(X)
            X_scaled = pd.DataFrame(X_scaled, columns=available_features, index=X.index)
        else:
            X_scaled = X
        
        # 5. 准备目标变量
        y = df_processed[target_col].values if target_col and target_col in df_processed.columns else None
        
        print(f"特征矩阵形状: {X_scaled.shape}")
        print(f"使用的特征列: {available_features}")
        if y is not None:
            print(f"目标变量形状: {y.shape}")
            print(f"目标变量统计: 均值={np.mean(y):.2f}, 标准差={np.std(y):.2f}")
        
        return X_scaled, y, df_processed

class SalesPredictor:
    """销售预测基础类"""
    
    def __init__(self, model_name):
        self.model_name = model_name
        self.model = None
        self.preprocessor = DataPreprocessor()
        self.is_trained = False
        
    def train(self, X, y):
        """训练模型"""
        raise NotImplementedError("子类必须实现train方法")
        
    def predict(self, X):
        """预测"""
        raise NotImplementedError("子类必须实现predict方法")
        
    def evaluate(self, X_test, y_test):
        """评估模型"""
        if not self.is_trained:
            raise ValueError("模型尚未训练")
            
        y_pred = self.predict(X_test)
        
        mae = mean_absolute_error(y_test, y_pred)
        mse = mean_squared_error(y_test, y_pred)
        rmse = np.sqrt(mse)
        r2 = r2_score(y_test, y_pred)
        
        return {
            'MAE': mae,
            'MSE': mse,
            'RMSE': rmse,
            'R2': r2
        }

class LinearRegressionPredictor(SalesPredictor):
    """线性回归预测器"""
    
    def __init__(self):
        super().__init__("LinearRegression")
        self.model = LinearRegression()
        
    def train(self, X, y):
        """训练线性回归模型"""
        self.model.fit(X, y)
        self.is_trained = True
        
    def predict(self, X):
        """预测"""
        if not self.is_trained:
            raise ValueError("模型尚未训练")
        return self.model.predict(X)

class RidgeRegressionPredictor(SalesPredictor):
    """岭回归预测器"""
    
    def __init__(self, alpha=1.0):
        super().__init__("RidgeRegression")
        self.model = Ridge(alpha=alpha)
        
    def train(self, X, y):
        """训练岭回归模型"""
        self.model.fit(X, y)
        self.is_trained = True
        
    def predict(self, X):
        """预测"""
        if not self.is_trained:
            raise ValueError("模型尚未训练")
        return self.model.predict(X)

class RandomForestPredictor(SalesPredictor):
    """随机森林预测器"""
    
    def __init__(self, n_estimators=100, random_state=42):
        super().__init__("RandomForest")
        self.model = RandomForestRegressor(
            n_estimators=n_estimators,
            random_state=random_state,
            n_jobs=-1
        )
        
    def train(self, X, y):
        """训练随机森林模型"""
        self.model.fit(X, y)
        self.is_trained = True
        
    def predict(self, X):
        """预测"""
        if not self.is_trained:
            raise ValueError("模型尚未训练")
        return self.model.predict(X)
        
    def feature_importance(self, feature_names):
        """获取特征重要性"""
        if not self.is_trained:
            raise ValueError("模型尚未训练")
        return dict(zip(feature_names, self.model.feature_importances_))

class GradientBoostingPredictor(SalesPredictor):
    """梯度提升预测器"""
    
    def __init__(self, n_estimators=100, learning_rate=0.1, random_state=42):
        super().__init__("GradientBoosting")
        self.model = GradientBoostingRegressor(
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            random_state=random_state
        )
        
    def train(self, X, y):
        """训练梯度提升模型"""
        self.model.fit(X, y)
        self.is_trained = True
        
    def predict(self, X):
        """预测"""
        if not self.is_trained:
            raise ValueError("模型尚未训练")
        return self.model.predict(X)

class TimeSeriesPredictor(SalesPredictor):
    """时间序列预测器（基于移动平均和趋势）"""
    
    def __init__(self, window=7):
        super().__init__("TimeSeries")
        self.window = window
        self.trend_params = None
        self.seasonal_pattern = None
        
    def train(self, X, y):
        """训练时间序列模型"""
        # 提取时间信息
        if 'days_from_start' in X.columns:
            time_index = X['days_from_start'].values
        else:
            time_index = np.arange(len(y))
            
        # 计算趋势
        if len(time_index) > 1:
            # 线性趋势拟合
            A = np.vstack([time_index, np.ones(len(time_index))]).T
            slope, intercept = np.linalg.lstsq(A, y, rcond=None)[0]
            self.trend_params = {'slope': slope, 'intercept': intercept}
        
        # 计算周期性模式
        if 'dayofweek' in X.columns:
            df_temp = pd.DataFrame({'y': y, 'dayofweek': X['dayofweek']})
            self.seasonal_pattern = df_temp.groupby('dayofweek')['y'].mean().values
        
        self.is_trained = True
        
    def predict(self, X):
        """预测"""
        if not self.is_trained:
            raise ValueError("模型尚未训练")
            
        predictions = np.zeros(len(X))
        
        # 基于趋势的预测
        if self.trend_params and 'days_from_start' in X.columns:
            time_index = X['days_from_start'].values
            trend_pred = self.trend_params['slope'] * time_index + self.trend_params['intercept']
            predictions += trend_pred
        
        # 基于周期性的预测
        if self.seasonal_pattern is not None and 'dayofweek' in X.columns:
            day_indices = X['dayofweek'].values.astype(int)
            seasonal_adjustment = self.seasonal_pattern[day_indices % 7]
            predictions += seasonal_adjustment
        
        # 确保预测值非负
        predictions = np.maximum(predictions, 0)
        
        return predictions